import argparse
import ast
import os
import tarfile
import zipfile
from pathlib import Path

import node_visitor
from detect import detect_dangercalls
from file_types_detect import get_executable_files
from name_detect import name_detect
from utils.logger import get_logger
from utils.utils import get_request, get_pypi_package_url, get_pypi_homepage_url, \
    get_repository_url, get_zipped_py_files

logger = get_logger()


def find_all_python_files(filepath):
    for root, dirs, files in os.walk(filepath):
        for file in files:
            if file.endswith('.py'):
                fullname = os.path.join(root, file)
                yield fullname


def get_python_file_list(packagename):
    python_file_list = []
    for python_file in find_all_python_files(packagename):
        # print(python_file)
        python_file_list.append(python_file)
    return python_file_list


def detect_package(python_file_list):
    danger_file_list = []
    for python_file in python_file_list:
        # print(python_file)
        try:
            file = open(python_file, 'rb')  # 以二进制只读方式读取Python文件
            source = file.read()
            node = ast.parse(source)  # 将源代码转化为抽象语法树
            visitor = node_visitor.NodeVisitor(python_file)
            visitor.generic_visit(node)  # 使用NodeVisitor类遍历抽象语法树
            dangercalls = visitor.dangercalls
            detect_result = detect_dangercalls(dangercalls)
            if detect_result:
                danger_file_list.append(python_file)
        except Exception as e:
            continue
    return danger_file_list


def detect_web(web, libraries_url, single=True):
    libraries_response = get_request(libraries_url)
    if not libraries_response:
        print('Package %s does not exist!!!' % web)
        logger.error('Package %s gets libraries response error!' % web)
        if single:
            exit(1)
        return False, []
    libraries_text = libraries_response.text
    pypi_homepage_url = get_pypi_homepage_url(libraries_text)
    pypi_url = pypi_homepage_url + '#files'
    repository_url = get_repository_url(libraries_text)
    pypi_response = get_request(pypi_url)
    if not pypi_response:
        print('Cannot get package %s on PyPI!!!' % web)
        logger.error('Package %s gets PyPI response error!' % web)
        if single:
            exit(1)
        return False, []
    href, file_name = get_pypi_package_url(pypi_response.text)
    if not href:
        print('Cannot get package %s on PyPI!!!' % web)
        logger.error('Package %s gets package url error!' % web)
        if single:
            exit(1)
        return False, []
    file_response = get_request(href)
    if not file_response:
        print('Cannot get package %s on PyPI!!!' % web)
        logger.error('Package %s gets package error!' % web)
        if single:
            exit(1)
        return False, []
    local_file = open('packages/' + file_name, 'wb')
    local_file.write(file_response.content)
    # if file_name.endswith('.zip'):
    #     zip_file=zipfile.ZipFile(file_response.content)
    #     zip_list=zip_file.namelist()
    # elif file_name.endswith('.tar.gz'):
    #     tar_file=tarfile.TarFile(file_response.content)
    if file_name.endswith('.zip'):
        zip_file = zipfile.ZipFile('packages/' + file_name)
        zip_list = zip_file.namelist()
        for file in zip_list:
            if os.path.splitext(file)[1] == ".py":
                zip_file.extract(file, 'packages')
        zip_file.close()
    elif file_name.endswith('.tar.gz'):
        # tar_file = tarfile.open(fileobj=file_response.content)
        tar_file = tarfile.open('packages/' + file_name)
        tar_file.extractall(members=get_zipped_py_files(tar_file), path='packages')
        tar_file.close()
    executable_file_list = get_executable_files('packages/' + file_name)
    if executable_file_list:
        logger.warning('Package %s has executable file(s): %s' % (web, str(executable_file_list)))
        if len(executable_file_list) == 1:
            print('Package %s has 1 executable file!!!' % web)
            print('The executable file is %s!' % executable_file_list[0])
        else:
            print('Package %s has %s executable files!!!' % web, str(len(executable_file_list)))
            print('There are the following files:')
            for executable_file in executable_file_list:
                print(executable_file)
    python_file_list = get_python_file_list('packages/' + file_name)
    danger_file_list = detect_package(python_file_list)
    return True, danger_file_list


def main():
    parser = argparse.ArgumentParser(description='Pypi package detector')
    parser.add_argument('-f', '--file')
    parser.add_argument('-p', '--package')
    parser.add_argument('-w', '--web')
    parser.add_argument('-v', '--version')
    parser.add_argument('-m', '--multiple')
    parser.add_argument('-l', '--list')
    args = parser.parse_args()
    filename = args.file
    packagename = args.package
    web = args.web
    version = args.version
    multiple = args.multiple
    package_list_file = args.list
    mode_num = 0
    if filename:
        mode_num += 1
    if packagename:
        mode_num += 1
    if web:
        mode_num += 1
    if multiple:
        mode_num += 1
    if package_list_file:
        mode_num += 1
    if mode_num > 1:
        print('Can only choose one mode!!!')
        logger.error('User chooses more than 1 mode!')
        exit(1)
    if mode_num == 0:
        print('Must choose one mode!!!')
        logger.error('User does not choose detection mode!')
        exit(1)
    if filename:
        file = None
        try:
            file = open(filename, 'r')
        except Exception as e:
            print('File does not exist!!!')
            logger.error('Open file %s error:%s!' % filename, e)
            exit(1)
        if not filename.endswith('.py'):
            print('File is not python file!!!')
            logger.error('File %s is not python file!' % filename)
            exit(1)
        source = file.read()
        node = ast.parse(source)
        visitor = node_visitor.NodeVisitor(filename)
        visitor.generic_visit(node)
        dangercalls = visitor.dangercalls
        detect_result = detect_dangercalls(dangercalls)
        if detect_result:
            print('File %s is dangerous!!!' % filename)
            logger.critical('Find dangerous file %s!' % filename)
            exit(0)
        else:
            print('File %s is not dangerous!' % filename)
            exit(0)
    if packagename:
        package = Path(packagename)
        if not package.exists():
            print('Folder does not exist!!!')
            logger.error('Folder %s does not exist!' % packagename)
            exit(1)
        executable_file_list = get_executable_files(packagename)
        if executable_file_list:
            logger.warning('Package %s has executable file(s): %s' % (packagename, str(executable_file_list)))
            if len(executable_file_list) == 1:
                print('Package %s has 1 executable file!!!' % packagename)
                print('The executable file is %s!' % executable_file_list[0])
            else:
                print('Package %s has %s executable files!!!' % packagename, str(len(executable_file_list)))
                print('There are the following files:')
                for executable_file in executable_file_list:
                    print(executable_file)
        python_file_list = get_python_file_list(packagename)
        if not python_file_list:
            print('Folder does not have python files!!!')
            logger.error('Folder %s has not have python files!' % packagename)
            exit(1)
        danger_file_list = detect_package(python_file_list)
        if danger_file_list:
            logger.critical('Package %s has dangerous file(s): %s' % (packagename, str(danger_file_list)))
            if len(danger_file_list) == 1:
                print('Package %s has 1 dangerous file!!!' % packagename)
                print('The dangerous file is %s!' % danger_file_list[0])
            else:
                print('Package %s has %s dangerous files!!!' % (packagename, str(len(danger_file_list))))
                print('There are the following files:')
                for danger_file in danger_file_list:
                    print(danger_file)
            exit(0)
        else:
            print('Not find dangerous file! The package is not dangerous!')
            exit(0)
    if multiple:
        multiple_package = Path(multiple)
        if not multiple_package.exists():
            print('Folder does not exist!!!')
            logger.error('Folder %s does not exist!' % multiple)
            exit(1)
        # package_list = os.listdir(multiple)
        package_list = []
        folders = []
        for root, dirs, files in os.walk(multiple):
            folders = dirs
            break
        for folder in folders:
            package = os.path.join(multiple, folder)
            package_list.append(package)
        if not package_list:
            print('Folder does not have packages!!!')
            logger.error('Folder %s does not have packages!' % multiple)
            exit(1)
        print('Detection starts. %s packages will be detected in all!' % str(len(package_list)))
        logger.info(
            'Detection on folder %s starts. %s packages will be detected in all!' % (multiple, str(len(package_list))))
        detected_package_list = []
        danger_package_list = []
        for package in package_list:
            print('')
            executable_file_list = get_executable_files(package)
            if executable_file_list:
                logger.warning('Package %s has executable file(s): %s' % (package, str(executable_file_list)))
                # danger_package_list.append(package)
                if len(executable_file_list) == 1:
                    print('Package %s has 1 executable file!!!' % package)
                    print('The executable file is %s!' % executable_file_list[0])
                else:
                    print('Package %s has %s executable files!!!' % (package, str(len(executable_file_list))))
                    print('There are the following files:')
                    for executable_file in executable_file_list:
                        print(executable_file)
            python_file_list = get_python_file_list(package)
            if not python_file_list:
                print('Package %s does not have python files!!!' % package)
                logger.error('Package %s does not have python files!!!' % package)
                continue
            # if len(python_file_list)>1000:
            #     continue
            danger_file_list = []
            # try:
            #     danger_file_list = detect_package(python_file_list)
            # except Exception as e:
            #     print('Error in detecting package %s!!!' % package)
            #     logger.error('Error in detecting package %s with error %s!' % (package, e))
            #     continue
            danger_file_list = detect_package(python_file_list)
            detected_package_list.append(package)
            if danger_file_list:
                logger.critical('Package %s has dangerous file(s): %s' % (package, str(danger_file_list)))
                danger_package_list.append(package)
                if len(danger_file_list) == 1:
                    print('Package %s has 1 dangerous file!!!' % package)
                    print('The dangerous file is %s!' % danger_file_list[0])
                else:
                    print('Package %s has %s dangerous files!!!' % (package, str(len(danger_file_list))))
                    print('There are the following files:')
                    for danger_file in danger_file_list:
                        print(danger_file)
            else:
                print('Not find dangerous file! Package %s is not dangerous!' % package)
        print('')
        print('Detection finished! %s packages have been correctly detected!' % str(len(detected_package_list)))
        logger.info('Detection on folder %s finished! %s packages have been correctly detected!' % (
            multiple, str(len(detected_package_list))))
        if len(danger_package_list) > 0:
            logger.critical('Folder %s has dangerous package(s): %s' % (multiple, str(danger_package_list)))
            if len(danger_package_list) == 1:
                print('1 dangerous package %s has been found!!!' % danger_package_list[0])
            elif len(danger_package_list) > 1:
                print('%s dangerous packages have been found!!!' % len(danger_package_list))
                print('There are the following packages:')
                for danger_package in danger_package_list:
                    print(danger_package)
        else:
            print('Have not found dangerous packages!')
    if web:
        typo_project_list = name_detect(web)
        if len(typo_project_list) > 0:
            logger.warning('Package %s may have typosquatting with package(s): %s!' % (web, str(typo_project_list)))
            if len(typo_project_list) == 1:
                print('Package %s may have typosquatting with package %s!!!' % (web, typo_project_list[0]))
            elif len(typo_project_list) > 1:
                print('Package %s may have typosquatting with %s packages!!!' % (web, len(typo_project_list)))
                print('There are the following packages:')
                for typo_project in typo_project_list:
                    print(typo_project)
        if version:
            libraries_url = 'https://libraries.io/pypi/' + web + '/' + version
        else:
            libraries_url = 'https://libraries.io/pypi/' + web
        flag, danger_file_list = detect_web(web, libraries_url)
        if len(danger_file_list) > 0:
            logger.critical('Package %s has dangerous file(s): %s' % (packagename, str(danger_file_list)))
            if len(danger_file_list) == 1:
                print('Package %s has 1 dangerous package!!!' % web)
                print('The dangerous file is %s!' % danger_file_list[0])
            else:
                print('Package %s has %s dangerous packages!!!' % (web, str(len(danger_file_list))))
                print('There are the following files:')
                for danger_file in danger_file_list:
                    print(danger_file)
            exit(0)
        else:
            print('Not find dangerous file! The package is not dangerous!')
            exit(0)
    if package_list_file:
        file = None
        try:
            file = open(package_list_file, 'r', encoding='utf-8')
        except Exception:
            print('File does not exist!!!')
            exit(1)
        lines = file.readlines()
        danger_package_dict = {}
        detected_package_list = []
        print('Detection starts!')
        logger.info('Detection starts!')
        for i in range(len(lines)):
            print('')
            line_split = lines[i].split(' ')
            for word in line_split:
                if not word:
                    line_split.remove(word)
            packagename = ''
            libraries_url = ''
            version = ''
            if len(line_split) == 1:
                packagename = line_split[0].strip()
                libraries_url = 'https://libraries.io/pypi/' + packagename
            elif len(line_split) == 2:
                packagename = line_split[0].strip()
                version = line_split[1].strip()
                libraries_url = 'https://libraries.io/pypi/' + packagename + '/' + version
            else:
                print('File format is wrong in line %s!!!' % str(i + 1))
                continue
            typo_project_list = name_detect(packagename)
            if typo_project_list:
                logger.warning(
                    'Package %s may have typosquatting with package(s): %s!' % (packagename, str(typo_project_list)))
                danger_package_dict[packagename] = version
                if len(typo_project_list) == 1:
                    print('Package %s may have typosquatting with package %s!!!' % (packagename, typo_project_list[0]))
                elif len(typo_project_list) > 1:
                    print(
                        'Package %s may have typosquatting with %s packages!!!' % (packagename, len(typo_project_list)))
                    print('There are the following packages:')
                    for typo_project in typo_project_list:
                        print(typo_project)
            flag, danger_file_list = detect_web(packagename, libraries_url, single=False)
            if not flag:
                continue
            detected_package_list.append((packagename, version))
            if len(danger_file_list) > 0:
                logger.critical('Package %s has dangerous file(s): %s' % (packagename, str(danger_file_list)))
                if packagename not in danger_package_dict:
                    danger_package_dict[packagename] = version
                if len(danger_file_list) == 1:
                    if version:
                        print('Package %s version %s has 1 dangerous file!!!' % (packagename, version))
                    else:
                        print('Package %s has 1 dangerous file!!!' % packagename)
                    print('The dangerous file is %s!' % danger_file_list[0])
                else:
                    if version:
                        print('Package %s version %s has %s dangerous files!!!' % (
                            packagename, version, str(len(danger_file_list))))
                    else:
                        print('Package %s has %s dangerous files!!!' % (packagename, str(len(danger_file_list))))
                    print('There are the following files:')
                    for danger_file in danger_file_list:
                        print(danger_file)
            elif len(danger_file_list) == 0:
                if version:
                    print(
                        'Not find dangerous file! The package %s version %s is not dangerous!' % (packagename, version))
                else:
                    print('Not find dangerous file! The package %s is not dangerous!' % packagename)
        print('')
        print('Detection finished! %s packages have been correctly detected!' % str(len(detected_package_list)))
        logger.info('File %s detection finished! %s packages have been correctly detected!' % (
            package_list_file, str(len(detected_package_list))))
        if len(danger_package_dict) > 0:
            logger.critical(
                'File %s detection has dangerous package(s): %s' % (package_list_file, str(danger_package_dict)))
            if len(danger_package_dict) == 1:
                for danger_package, danger_version in danger_package_dict.items():
                    print('1 dangerous package %s version %s has been found!!!' % danger_package, danger_version)
            elif len(danger_package_dict) > 1:
                print('%s dangerous packages have been found!!!' % len(danger_package_dict))
                print('There are the following packages:')
                for danger_package, danger_version in danger_package_dict.items():
                    if danger_version:
                        print(danger_package, danger_version)
                    else:
                        print(danger_version)
        else:
            print('Have not found dangerous packages!')


main()
